# Copyright 2021-2023 @ Shenzhen Bay Laboratory &
#                       Peking University &
#                       Huawei Technologies Co., Ltd
#
# This code is a part of MindSPONGE:
# MindSpore Simulation Package tOwards Next Generation molecular modelling.
#
# MindSPONGE is open-source software based on the AI-framework:
# MindSpore (https://www.mindspore.cn/)
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ============================================================================
"""MindSPONGE"""

import time
from distutils.version import LooseVersion
import logging


def _pytorch_version_check():
    """
    Do the PyTorch version check. If PyTorch can not be imported, it will raise ImportError. 
    If its version is not compatible with current version, it will print a warning.

    Raises:
        ImportError: If PyTorch can not be imported.
    """
    try:
        import torch
        # 配置logging
        logging.basicConfig(level=logging.INFO)
        logger = logging.getLogger(__name__)
    except ImportError as ie:
        raise ImportError("Can not find PyTorch in current environment. Please install "
                         "PyTorch before using this package, by following "
                         "the instruction at https://pytorch.org/get-started/locally/") from ie

    torch_version = torch.__version__
    required_torch_version = '2.0.0'  # 设置所需的最低PyTorch版本
    logger.info(f"Current PyTorch version is {torch_version}")
    
    if LooseVersion(torch_version) < LooseVersion(required_torch_version):
        logger.warning(f"Current version of PyTorch is not compatible with this package. "
                      f"Some functions might not work or even raise error. Please install PyTorch "
                      f"version >= {required_torch_version}. "
                      f"For more details about dependency setting, please check "
                      f"the instructions at PyTorch official website https://pytorch.org/get-started/locally/")
        
        warning_countdown = 3
        for i in range(warning_countdown, 0, -1):
            logger.warning(f"Please pay attention to the above warning, countdown: {i}")
            time.sleep(1)

_pytorch_version_check()

# pylint: disable=wrong-import-position
from .system import get_molecule, Molecule, Protein, MoleculeFromMol2
from .potential import PotentialCell, ForceFieldBase, ForceField
from .optimizer import Updater, UpdaterMD, DynamicUpdater
from .core import Sponge, WithEnergyCell, SimulationCell, RunOneStepCell, AnalysisCell, WithForceCell
from .function.units import GLOBAL_UNITS, set_global_units, set_global_length_unit, set_global_energy_unit
